
import os, math, time, argparse, shutil
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from tqdm import tqdm

from gaussian_sr.models.gaussian_sr import GaussianSRModel
from dataset.sr_dataset import SRHRDataset
from losses import L1Loss
from utils.common import sample_hr_pixels, gather_pixels, psnr

def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--data_root", type=str, default="./data")
    p.add_argument("--save_dir", type=str, default="./runs/gaussian_sr")
    p.add_argument("--epochs", type=int, default=300)
    p.add_argument("--batch_size", type=int, default=8)
    p.add_argument("--num_workers", type=int, default=4)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--channels", type=int, default=64)
    p.add_argument("--blocks", type=int, default=8)
    p.add_argument("--gaussians", type=int, default=100)
    p.add_argument("--temperature", type=float, default=1.0)
    p.add_argument("--hr_crop", type=int, default=128)
    p.add_argument("--samples", type=int, default=2304)
    p.add_argument("--scale_min", type=float, default=1.1)
    p.add_argument("--scale_max", type=float, default=4.0)
    p.add_argument("--clip_grad", type=float, default=1.0)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--eval_only", action="store_true")
    p.add_argument("--ckpt", type=str, default="")
    return p.parse_args()

def seed_everything(seed):
    import random, numpy as np, torch
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

def save_ckpt(state, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save(state, path)

def load_ckpt(model, opt, path):
    ck = torch.load(path, map_location="cpu")
    model.load_state_dict(ck["model"], strict=True)
    if opt is not None and "opt" in ck and ck["opt"] is not None:
        opt.load_state_dict(ck["opt"])
    return ck.get("epoch", 0), ck.get("best_psnr", None)

@torch.no_grad()
def evaluate(model, loader, device, scales=(2,3,4)):
    model.eval()
    total_psnr = 0.0
    n = 0
    for batch in loader:
        hr = batch["hr"].to(device)
        for s in scales:
            lr = torch.nn.functional.interpolate(hr, scale_factor=1.0/s, mode="bicubic", align_corners=False)
            B, _, Hh, Wh = hr.shape
            sr, _, _ = model(lr, (Hh, Wh), s, train_mode=False)
            total_psnr += psnr(sr.clamp(0,1), hr).item()
            n += 1
    return total_psnr / max(n,1)

def main():
    args = get_args()
    seed_everything(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    train_set = SRHRDataset(args.data_root, split="train", hr_crop=args.hr_crop,
                            scale_min=args.scale_min, scale_max=args.scale_max)
    val_set   = SRHRDataset(args.data_root, split="val",   hr_crop=args.hr_crop,
                            scale_min=2.0, scale_max=4.0)
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
                              num_workers=args.num_workers, pin_memory=True, drop_last=True)
    val_loader   = DataLoader(val_set, batch_size=1, shuffle=False,
                              num_workers=args.num_workers, pin_memory=True)

    model = GaussianSRModel(in_ch=3, channels=args.channels, num_blocks=args.blocks,
                            num_gaussians=args.gaussians, temperature=args.temperature,
                            use_alpha_sum=True).to(device)

    opt = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999))
    sched = torch.optim.lr_scheduler.StepLR(opt, step_size=200, gamma=0.5)
    scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
    criterion = L1Loss()

    start_epoch = 0
    best_psnr = -1e9
    os.makedirs(args.save_dir, exist_ok=True)

    if args.ckpt:
        start_epoch, best_psnr = load_ckpt(model, opt if not args.eval_only else None, args.ckpt)
        print(f"Loaded checkpoint from {args.ckpt} (epoch={start_epoch}, best_psnr={best_psnr})")
        if args.eval_only:
            val_psnr = evaluate(model, val_loader, device)
            print(f"[EvalOnly] PSNR: {val_psnr:.3f} dB")
            return

    if args.eval_only:
        raise ValueError("Provide --ckpt for eval_only.")

    for epoch in range(start_epoch, args.epochs):
        model.train()
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}")
        avg_loss = 0.0
        for batch in pbar:
            lr = batch['lr'].to(device)
            hr = batch['hr'].to(device)
            scale = batch['scale']
            Hh, Wh = batch['hr_size']
            if isinstance(scale, float) or isinstance(scale, int):
                s_val = float(scale)
                scale = torch.tensor([s_val]*lr.shape[0], device=device)
            else:
                scale = scale.to(device)
            Hh = int(Hh); Wh = int(Wh)

            opt.zero_grad(set_to_none=True)
            with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
                sr, _, _ = model(lr, (Hh, Wh), scale[0].item(), train_mode=True)

                ys, xs = sample_hr_pixels(Hh, Wh, N=args.samples, device=device)
                gt_pix = gather_pixels(hr, ys, xs)
                pr_pix = gather_pixels(sr, ys, xs)
                loss = criterion(pr_pix, gt_pix)

            scaler.scale(loss).backward()
            if args.clip_grad > 0:
                scaler.unscale_(opt)
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad)
            scaler.step(opt)
            scaler.update()

            avg_loss += loss.item()
            pbar.set_postfix(loss=f"{loss.item():.4f}")

        sched.step()
        avg_loss /= max(len(train_loader),1)

        val_psnr = evaluate(model, val_loader, device)
        print(f"Epoch {epoch+1}: train_loss={avg_loss:.4f}, val_psnr={val_psnr:.3f} dB")
        ckpt_path = os.path.join(args.save_dir, "last.ckpt")
        save_ckpt({"model": model.state_dict(), "opt": opt.state_dict(),
                   "epoch": epoch+1, "best_psnr": best_psnr}, ckpt_path)
        if val_psnr > best_psnr:
            best_psnr = val_psnr
            shutil.copy(ckpt_path, os.path.join(args.save_dir, "best.ckpt"))

        with torch.no_grad():
            sample = next(iter(val_loader))
            hr = sample["hr"].to(device)
            for s in [2.0, 3.0, 4.0]:
                lr = torch.nn.functional.interpolate(hr, scale_factor=1.0/s, mode="bicubic", align_corners=False)
                B, _, Hh, Wh = hr.shape
                sr_img, _, _ = model(lr, (Hh, Wh), s, train_mode=False)
                os.makedirs(os.path.join(args.save_dir, "preview"), exist_ok=True)
                save_path = os.path.join(args.save_dir, "preview", f"epoch{epoch+1}_x{s}.png")
                save_image(torch.clamp(sr_img, 0, 1), save_path)

if __name__ == "__main__":
    main()
